import numpy as np

from environ.env_5bits_2048 import Env_2048


class ENV_VEC(Env_2048):
    def __init__(self) -> None:
        super().__init__()

    def init_board(self, env_num):
        board = np.zeros((env_num, 16), dtype=np.uint16)
        penalty = np.zeros(env_num, dtype=np.int64)

        for i in range(env_num):
            index_a, index_b = np.random.choice(16, 2, replace=False)
            tile_a, tile_b = np.random.choice([1, 2], size=2, p=[0.9, 0.1], replace=True)
            board[i][index_a], board[i][index_b] = tile_a, tile_b
            penalty[i] = ((tile_a>>1) + (tile_b>>1))*4

        return board.reshape((env_num, 4, 4)), penalty
         

    def move_left(self, board:np.array):
        a = board[..., 0]
        b = board[..., 1]
        c = board[..., 2]
        d = board[..., 3]

        new_board = self.move_left_table[a, b, c, d]
        return new_board

    def move_right(self, board:np.array):
        a = board[..., 0]
        b = board[..., 1]
        c = board[..., 2]
        d = board[..., 3]

        new_board = self.move_right_table[a, b, c, d]
        return new_board

    def move_up(self, board:np.array):
        t_board = self.move_left(np.transpose(board, (0, 2, 1)))
        return np.transpose(t_board, (0, 2, 1))

    def move_down(self, board:np.array):
        t_board = self.move_right(np.transpose(board, (0, 2, 1)))
        return np.transpose(t_board, (0, 2, 1))


    def insert_tile(self, board:np.array):
        """
        single_board=np.array([[0, 0, 0, 0],
                               [0, 1, 0, 0],
                               [0, 0, 0, 0], 
                               [0, 0, 1, 0]], dtype=np.uint16) 
        """
        env_num = len(board)
        true_board = (board == 0)  # 必然每个环境都有空格
        empty_num = np.sum(true_board, axis=(1,2))  # 空格数的范围：1~15
        # 指定插入空格的索引号：1~15
        index = np.array(np.random.random_sample(env_num)*empty_num, dtype=np.int16)+1

        tile = np.random.choice([1, 2], size=env_num, p=[0.9, 0.1], replace=True)

        for i in range(4):
            for j in range(4):
                index -= true_board[:, i, j]
                tile_flag = (index == 0)
                board[:, i, j] += np.array(tile_flag*tile, dtype=np.uint16)
                index -= tile_flag*100  # 如果不设置该语句，那么在插入的位置的下一位有数值会对其进行再次插入加和

        penalty = (tile>>1)*4
        return board, penalty


    def move(self, board:np.array, move_left_index, move_right_index, move_up_index, move_down_index):
        if(len(move_left_index)!=0):
            left_board = self.move_left(board[move_left_index])
            board[move_left_index] = left_board

        if(len(move_right_index)!=0):
            right_board = self.move_right(board[move_right_index])
            board[move_right_index] = right_board

        if(len(move_up_index)!=0):
            up_board = self.move_up(board[move_up_index])
            board[move_up_index] = up_board

        if(len(move_down_index)!=0):
            down_board = self.move_down(board[move_down_index])
            board[move_down_index] = down_board

        board, penalty = self.insert_tile(board)
        return board, penalty


    def part_reset(self, board:np.array, penalty:np.array):
        unava_bool_index = self.get_unavailable_index(board)

        if(np.sum(unava_bool_index)==0):
            return board, penalty 

        score = self.get_score(board[unava_bool_index], penalty[unava_bool_index])
        max_rank = self.get_max_rank(board[unava_bool_index])
        print(score)
        print(max_rank)

        part_board, part_penalty = self.init_board(np.sum(unava_bool_index))
        board[unava_bool_index] = part_board
        penalty[unava_bool_index] = part_penalty
        return board, penalty


    def get_unavailable_index(self, board:np.array)->np.array(np.bool_):
        ava_array = self.get_available(board)
        unava_bool_index = np.sum(ava_array, axis=0)==0
        return unava_bool_index

    
    def get_available(self, board:np.array):
        left_board = self.move_left(board)
        right_board = self.move_right(board)
        up_board = self.move_up(board)
        down_board = self.move_down(board)

        ava_left = np.sum(left_board != board, axis=(1, 2))
        ava_right = np.sum(right_board != board, axis=(1, 2))
        ava_up = np.sum(up_board != board, axis=(1, 2))
        ava_down = np.sum(down_board != board, axis=(1, 2))
        ava = np.array([ava_left, ava_right, ava_up, ava_down], dtype=np.bool_)
        return ava


    def count_empty(self, board:np.array):
        return np.sum(board == 0, axis=(1,2))

    def get_max_rank(self, board):
        return np.max(board, axis=(1,2))

    def get_score(self, board:np.asarray, penalty:np.asarray):
        a = board[..., 0]
        b = board[..., 1]
        c = board[..., 2]
        d = board[..., 3]

        score = np.sum( self.score_table[a, b, c, d] )
        return score-penalty


if __name__ == '__main__':
    env = ENV_VEC()
    env_num = 100
    board, penalty = env.init_board(env_num)

    for t in range(int(1e6)):
        print(t)
        # print(board)
        left = []
        right = []
        up = []
        down = []
        ava_array = env.get_available(board)
        for i in range(env_num):
            cc = np.random.choice(np.where(ava_array[:, i])[0])
            if(cc == 0):
                left.append(i)
            if(cc == 1):
                right.append(i)
            if(cc==2):
                up.append(i)
            if(cc==3):
                down.append(i)
        board, p = env.move(board, np.array(left), np.array(right), np.array(up), np.array(down))
        penalty += p
        board, penalty = env.part_reset(board, penalty)


